# encoding: utf-8
# !/usr/local/bin/python3

"""
@author: Gao Shuo
@contact: dorothy400@163.com
@software: PyCharm
@file: test.py
@time: 2020/3/8 10:04

input:
output:
"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import tensorflow as tf
from tensorflow.python.platform import app
import Model
# from quiz_recog.train_model import Model
import os
import numpy as np
FLAGS = app.flags.FLAGS
app.flags.DEFINE_string("test_path", "./data/test1", "the path of test images")
app.flags.DEFINE_string("log_dir", "./model/model.ckpt", "checkpoint dir")


# 建立test图片列表
def _test_image_list(file_path):
    img_list = []
    for image in os.listdir(file_path):
        img_list.append(file_path + '/' + image)
    return img_list


def images_read(file_list):
    """
    读取汉字的图片
    :param file_list: 文件列表名
    :return: 每张图片的张量
    """
    # 1、构建文件队列
    file_enq = tf.train.string_input_producer(file_list)
    # 2、构建图片阅读器
    reader = tf.WholeFileReader()
    key, value = reader.read(file_enq)
    # 3、解码图片
    image = tf.image.decode_jpeg(value, channels=3)
    # 4、修改图片大小
    image_resize = tf.image.resize_images(image, [224, 224])
    image_resize.set_shape([224, 224, 3])
    # 5、读取多张图片，批处理
    image_batch = tf.train.batch(
        [image_resize],
        batch_size=1,
        num_threads=1,
        capacity=10)
    return image_batch


def main():
    file_path = FLAGS.test_path
    file_list = _test_image_list(file_path)
    # 得到图片读取内容
    image_batch = images_read(file_list=file_list)
    logits, end_points = Model.DenseNet(image_batch, num_class=num_classes, is_training=False,
                                        dropout_keep_prob=1, scope='DenseNet')
    predict = tf.nn.softmax(logits)
    x = tf.placeholder(tf.float32, [None, 224, 224, 3])
    saver = tf.train.Saver()

    # 开启会话
    with tf.Session() as sess:
        init_op = tf.group(
            tf.global_variables_initializer(),
            tf.local_variables_initializer())
        sess.run(init_op)
        saver.restore(sess, FLAGS.log_dir)
        # 构建线程管理器
        coord = tf.train.Coordinator()
        # 开启子线程
        threads = tf.train.start_queue_runners(sess=sess, coord=coord)
        # 主线程处理数据
        images = sess.run([image_batch])[0]
        pre = sess.run(predict, feed_dict={x: images})
        max_ind = np.argmax(pre)
        # 关闭线程管理器
        coord.request_stop()
        coord.join(threads=threads)
        print(max_ind)


if __name__ == '__main__':
    num_classes = 100
    main()
